import json

from click import prompt
from pyasn1_modules.rfc5652 import otherKeyAttributesMap
from sympy.crypto import padded_key

from index_admin import vector_load
from restful import RequestData, MessageItem
from datetime import datetime
from time import sleep
import logging
from  api_doubao import doubao_chat,doubao_chat_stream
from api_deepseek import deepseek_stream, deepseek_chat
from concurrent.futures import ThreadPoolExecutor

logging.basicConfig(
    level=logging.INFO,  # 设置日志级别
    filename='llm_knowledge_base.log',  # 指定日志文件名
    filemode='a',  # 'a'表示追加模式，'w'表示覆盖模式
    format='%(asctime)s - %(levelname)s - %(message)s'  # 自定义日志格式
)

PDF_LINK= f'http://132.10.10.121:8100/download_pdf/'

def mutil_function(function,input_datas):
    with ThreadPoolExecutor() as executor:
        results = executor.map(function, input_datas)
    return results

class ChangeName:
    @staticmethod
    def replace_str(line):
        new_line = (line.replace("：", "_").
         replace("、", "_").
         replace("（", "_").
         replace("）", "_").
         replace("‘", "_").
         replace("’", "_").
         replace("“", "_").
         replace("”", "_").
         replace("，", "_").replace("】", "_").replace("【", "_").replace("@","_").
        replace(" ","").replace("《","").replace("》",""))
        return new_line
    @staticmethod
    def replace_start(line:str):
        while line.startswith("_"):
            line = line[1:]
        return line[:-1]

    @staticmethod
    def sum_len(filename):
        s_name = filename.split("_")
        ns_name = [ChangeName.replace_str(item) for item in s_name]
        reversed_list = ns_name[::-1]
        temp = ""
        for idx , item in enumerate(reversed_list):
            if idx == len(reversed_list)-1:
                return temp
            else:
                if len(item)>70:
                    return item
                elif len(item) + len(temp) >90:
                    return temp
                else:
                    temp =  item+"_" + temp


class Tools:
    @staticmethod
    def remove_duplicate_rows(lst):
        first_elements = set()
        result = []
        for row in lst:
            if row[0] not in first_elements:
                result.append(row)
                first_elements.add(row[0])
        return result

    @staticmethod
    def process_response(response):
        return json.loads(response)

    @staticmethod
    def org_response(response):
        return response

    @staticmethod
    def process_err(function, res_function, max_rep, prompt):
        rep_e = 0
        while rep_e < max_rep:
            try:
                rt = function(prompt)
                right_rt = res_function(rt)
                return right_rt
            except Exception as e:
                if max_rep != 1:
                    print(f"捕获到错误: {e}, 正在尝试重新执行...")
                rep_e += 1
                sleep(2)
                if rep_e == max_rep:
                    if max_rep != 1:
                        print("已达到最大重试次数，操作失败。")
                    return "抽取/生成失败"

    @staticmethod
    def create_prompt_out_search_one(question):
#         prompt = '你的角色是一个AI助手，\
# 主要工作内容是处理用户查询，生成可用于向量库搜索匹配的语料。\n\
# 工作流程/工作任务如下：\n\
# 分析用户查询：理解用户输入的问题，提取关键信息。\n\
# 构建搜索语料：根据关键信息，构建适用于向量库搜索的语句或短语。\n\
# 优化搜索语句：调整语句结构，确保其能有效匹配向量库中的相关内容。\n\
# 以下是一个案例：\n\
# 用户查询：通达系转运中心数量。\n\
# 搜索语料：["通达系转运中心数量","中通快递转运中心数量","圆通速递转运中心数量","申通快递转运中心数量","韵达快递转运中心数量"]\n\
# 注意事项：\n\
# 确保搜索语料与用户查询紧密相关。\n\
# 使用简洁、明确的语句，避免模糊不清的表述。\n\
# 考虑不同关键词的组合，以提高搜索的准确性和覆盖面。\n\
# 用户查询为:'+question+"\n"+'输出格式为json,字段用英文双引号引起来,XXX代表你回答的内容：{"搜索语料":["XXX","XXX"]}'
#         prompt = "现在有一套基于大模型和RAG的智能问答系统，系统将在向量数据库中搜索关键词并将返回结果作为大模型的语料，为了使RAG系统召回的结果更加全面准确，\
# 请你根据用户的问题给出用于RAG系统的检索词，请不要输出检索词以外的内容。用户的问题如下：\n"+"'''\n"+question+"\n'''"+'\n输出格式为json,字段用英文双引号引起来,XXX代表你回答的内容：{"检索词":["XXX","XXX"]}'
#         prompt = ("现在有一套基于大模型和RAG的智能问答系统，系统将在向量数据库中搜索关键词并将返回结果作为大模型的语料。为了使RAG系统召回的结果更加全面准确，请改写用户的问题，改写规则如下：\n\
# 当用户输入的问题中存在无法通过向量化语义识别出的简写和代称时，请给出使用同义词进行替换并输出原问题和词语替换后的问题， 例如\
# 通达系转运中心数量需要把问题改写，改写后的问题为：[通达系转运中心数量, 中通转运中心数量, 圆通转运中心数量, 申通转运中心数量, 韵达转运中心数量]。但\
# 顺丰发展历程则不需要改写，改写后的问题和原问题一样:[顺丰发展历程]。根据要求和对应的案例，改写如下问题：\n\
# '''\n"+ question+"'''\n"+"输出格式为json,字段用英文双引号引起来,XXX代表你回答的内容：{\"改写后的问题\":[\"XXX\"]}")
#         prompt = "现在有一套基于大模型和RAG的智能问答系统，系统将在向量数据库中搜索关键词并将返回结果作为大模型的语料。为了使RAG系统召回的结果更加全面准确，请改写用户的问题，改写规则： 当用户输入的问题中存在无法通过向量化语义识别出的简写和代称时，请给出使用同义词进行替换并输出原问题和词语替换后的问题， 例如当用户输入“通达系转运中心数量”，输出为json list [ \"通达系转运中心数量\", \"中通转运中心数量\", \"圆通转运中心数量\", \"申通转运中心数量\", \"韵达转运中心数量\"] 请不要输出json list以外的内容。用户的问题如下: \n'''\n"+question+"\n'''"
        prompt = "现在有一套基于大模型和RAG的智能问答系统，系统将在向量数据库中搜索关键词并将返回结果作为大模型的语料。为了使RAG系统召回的结果更加全面准确，请改写用户的问题，改写规则： 当用户输入的问题中存在简写和代称时，请给出使用全称进行替换并输出原问题和词语替换后的问题， 例如当用户输入“通达系转运中心数量”，输出为json list [ \"通达系转运中心数量\", \"中通转运中心数量\", \"圆通转运中心数量\", \"申通转运中心数量\", \"韵达转运中心数量\"] 请不要输出json list以外的内容。用户的问题如下: \n'''\n" + question + "\n'''"
        return prompt

    @staticmethod
    def search_vdb(question ,k ,bot):
        vector_data = vector_load(bot)
        vector_search = vector_data.similarity_search_score(question, k=k)
        return vector_search

    @staticmethod
    def process_chroma_select_simple(question,k,bot):
        vector_search = Tools.search_vdb(question=question,k=k,bot=bot)
        comes = []
        content = []
        pdf = []
        score = []
        for doc in vector_search:
            # print(f"doc:{doc[1]},{doc[0]}")
            useful = doc[0].to_json()['kwargs']
            # print(f"useful:{useful}")
            useful_content = useful["page_content"] + "该文字片段来自于[" + useful["metadata"]["org_article_name"] + "]"
            content.append(useful_content)
            pdf_name = useful["metadata"]['org_article_name']
            if pdf_name in pdf:
                continue
            pdf.append(pdf_name)
            comes.append(PDF_LINK + pdf_name + '.pdf')
            score.append(str(round(float(doc[1]), 2)))
        print(f"{question}对应向量库检索的文章片段：{content}")
        return content, comes, score

    @staticmethod
    def process_same_comes(useful_pdf_name,useful_score):
        pdf_name = []
        res = []
        for a, b in zip(useful_pdf_name, useful_score):
            if a in pdf_name:
                continue
            pdf_name.append(a)
            temp = PDF_LINK +a+".pdf"+ "**" + b
            res.append(temp)
        return res

    @staticmethod
    #考虑如果用户的问题是多类别？？？
    def classify_question(question):
        type_question = ["要实际发生的信息", "要预测信息"]
        ut = datetime.now().strftime("%Y-%m-%d")
        prompt_one = "用户的问题是：" + question + "\n"

        # prompt_one = "用户的问题是：" + question + "用户所在时间是："+datetime.now().strftime("%Y-%m-%d") +"\n"
        prompt_two = "当前时间是" + datetime.now().strftime("%Y-%m-%d") + "，请先思考当前时间是何时，然后再分析用户的问题是" + "还是".join(type_question) + "\n"
        # prompt_two = "当前时间是" + datetime.now().strftime("%Y-%m-%d") + "，首先分析用户问题中涉及的时间，其次判断该时间与当前时间关系，如果用户问题中涉及的时间早于当前时间，属于客观事实问题，如果问题中询问的时间信息晚于当前时间，属于预测问题"+ "\n"+"例如用户询问1992年中国的GDP值，用户问题中的时间是1992年比当前时间"+datetime.now().strftime("%Y-%m-%d")+"早，那么则属于针对事实的提问，属于客观事实问题"
        # prompt_two = "请判断用户的问题类别，是属于寻求历史客观事实，还是对未来的预测"+"，例如；用户询问圆通2024年的利润，用户在时间是："+ut+"，问题中涉及的时间比用户所在时间早，对于用户来说，这是已经发生的历史，属于客观事实问题。"
        prompt_three = "输出格式为json，\n{\n\"问题类型\":\n\"理由\":\n}"
        prompt = prompt_one + prompt_two + prompt_three
        print(f"prompt：{prompt}")
        res = Tools.process_err(doubao_chat,Tools.process_response,3,prompt)
        return res

    @staticmethod
    def sub_mutil(function ,prompts,get_str):
        group_n = 20
        res = []
        if len(prompts)>group_n:
            for i in range(0,len(prompts),group_n):
                part_prompt = prompts[i:i+group_n]
                mf = mutil_function(function,part_prompt)
                mf_ =  [i for i in mf]
                res = res  + mf_
                sleep(1)
        else:
            ans_g = mutil_function(function, prompts)
            res = [i for i in ans_g]
        return_res = []
        for i in res:
            try:
                j_i = json.loads(i)
                temp_i = j_i[get_str]
            except Exception as e:
                print(f"err:{e}")
                temp_i = "大模型生成失败"
            return_res.extend(temp_i)
        return return_res



    @staticmethod
    def process_chroma_select_mutil(question,k,bot,function):
        vector_search = Tools.search_vdb(question=question,k=k,bot=bot)
        useful = [doc[0].to_json()['kwargs'] for doc in vector_search]
        content = [uf["page_content"] + "该文字片段来自于[" + uf["metadata"]["org_article_name"] + "]" for uf in useful]
        # print(f"content:{useful}")
        vt_score = [str(round(float(doc[1]), 2)) for doc in vector_search]
        pdf_name = [uf["metadata"]['org_article_name'].replace(".md","") for uf in useful]
        end_prompt = "\n请判断上述文字是否包含回答问题为\"" + question + "\"的答案,请注意，当文字片段中包含预测，而问题中没有时，不能视作可以回答该问题的片段，例如\n\
“\n\
问题：顺丰2024年利润是多少？\n\
文字片段：交通运输行业2022年投资策略：追寻疫情下的确定，重视拐点类标的：针对顺丰速运（顺丰控股的速运分部），在保守假设下，净利润为187亿元；在中性假设下，净利润为202亿元；在乐观假设下，净利润为243亿元。\n\
”\n\
只用输出是或者否，其他都不用输出"
        prompts = [i + end_prompt for i in content]
        print(f"批量判断文件块是否和问题相关")
        logging.info(f"批量判断文件块是否和问题相关")
        ans_g = mutil_function(function, prompts)
        temp = [i for i in ans_g]
        logging.info(f"文件块相关性结果：{temp}")
        try:
            temp_ans = [i.split("</think>")[1].strip() for i in temp]
        except:
            temp_ans = [i.strip() for i in temp]
        ans = [idx for idx, item in enumerate(temp_ans) if "否" not in item]
        useful_content = [item for idx, item in enumerate(content) if idx in ans]
        useful_score = [item for idx, item in enumerate(vt_score) if idx in ans]
        useful_pdf_name =  [item for idx, item in enumerate(pdf_name) if idx in ans]
        # print(f"useful_content:{useful_content}")
        # comes = [PDF_LINK +a+".pdf"+ "**" + b for a,b in zip(useful_pdf_name,useful_score)]
        comes = Tools.process_same_comes(useful_pdf_name,useful_score)
        return useful_content,comes




    @staticmethod
    def process_chroma_select_mutil_v1(question, k, bot, function):
        vector_search = Tools.search_vdb(question=question, k=k, bot=bot)
        useful = [doc[0].to_json()['kwargs'] for doc in vector_search]
        content = [uf["page_content"] + "该文字片段来自于[" + uf["metadata"]["org_article_name"] + "]" for uf in
                   useful]
        # print(f"content:{useful}")
        vt_score = [str(round(float(doc[1]), 2)) for doc in vector_search]
        pdf_name = [uf["metadata"]['org_article_name'].replace(".md", "") for uf in useful]

        print(f"判断问题类型")
        type_question = Tools.classify_question(question)
        print(f"用户问题的类别是：{type_question}")


        print(f"批量判断文件块是否和问题相关")
        logging.info(f"批量判断文件块是否和问题相关")
        prompt_one = "以上为从知识库中检索的片段信息以及文章出处，需要回答的问题是："+question +"\n"
        prompt_two = "当前时间是" + datetime.now().strftime("%Y-%m-%d") + "\n"
        prompt_three = "请记住用户的问题是"+type_question["问题类型"]+"，输出格式为json，\n\
{\n\
“片段是否可以回答问题”:是或者否\n\
“理由”:\n\
}"
        prompts = [text +"\n"+prompt_one + prompt_two + prompt_three for text in content]
        smr = Tools.sub_mutil(doubao_chat,prompts,"片段是否可以回答问题")
        logging.info(f"文件块相关性结果：{smr}")
        print(f"文件块相关性结果：{smr}")
        temp_ans = [i for i in smr]
        ans = [idx for idx, item in enumerate(temp_ans) if "是" in item]
        useful_content = [item for idx, item in enumerate(content) if idx in ans]
        useful_score = [item for idx, item in enumerate(vt_score) if idx in ans]
        useful_pdf_name = [item for idx, item in enumerate(pdf_name) if idx in ans]
        # print(f"useful_content:{useful_content}")
        # comes = [PDF_LINK +a+".pdf"+ "**" + b for a,b in zip(useful_pdf_name,useful_score)]
        comes = Tools.process_same_comes(useful_pdf_name, useful_score)
        return useful_content, comes


    @staticmethod
    def combine_ans_out_search_one(ans:list,questions:list,o_q:str):
        part_one ="原始问题为："+o_q+"对其改写后的问题和答案如下：\n"
        part_two = [f"改写后的问题：\n{i}\n改写后问题的答案：\n{j}\n" for i ,j in zip(questions,ans)]
        useful_part_two = " ".join(part_two)
        part_three = "根据改写后的问题及答案，整理归纳并回答原始问题"+o_q
        part_four = "要求答案忠实于原文，避免推断，逻辑合理，按由远及近的时间线进行回答，针对含有数字的答案，需二次核对校验，采用条列式结构输出"
        prompt = part_one + useful_part_two + part_three + part_four
        return prompt

    @staticmethod
    def combine_vb_ans(contents:list,question:str):
        part_end = "要求答案忠实于原文，避免推断，逻辑合理，按由远及近的时间线进行回答，针对含有数字的答案，需二次核对校验，采用条列式结构输出"
        combine_question = "'''" + "\n" + "\n\n".join(contents) + "\n'''" + "\n请根据上述材料回答：" + str(question) + part_end
        return combine_question



class Api:
    def __init__(self):
        self.tool = Tools

    def stream_chat_account(self,question):
        prompt = self.tool.create_prompt_account(question)
        try:
            ans = self.tool.process_err(function=doubao_chat,
                              res_function=self.tool.process_response,
                              max_rep=3,
                              prompt=prompt)
        except Exception as e:
            print(f"err:{e}")
            ans = {'回答': '大模型回答失败,请重新尝试','文件出处': '', '理由': ''}
        return  ans, prompt

    def chat_out_search(self,question):
        prompt = self.tool.create_prompt_out_search_one(question)
        print(f"prompt:{prompt}")
        try:
            ans = self.tool.process_err(function=doubao_chat,
                                        res_function=self.tool.process_response,
                                        max_rep=3,
                                        prompt=prompt)

        except Exception as e:
            print(f"err:{e}")
            ans=[question]
        new_ans = ans
        return new_ans

    def combine_ans_out_search(self,ans:list,questions:list,o_q:str):
        prompt = self.tool.combine_ans_out_search_one(ans,questions,o_q)
        print(f"prompt:{prompt}")
        ans = doubao_chat_stream(prompt)
        return ans




def chat_base(request: RequestData):
    #获取要使用的助手
    bot = request.bot
    #获取问题
    question = request.messages[-1].content
    # 构造含有本地知识库内容的问题请求
    api = Api()
    tool = Tools
    # deepseek_规范化管理模块:
    if bot == "Template_Warehouse":
        print("规范化管理模块")
        logging.info("规范化管理模块")
        contents,come,score =tool.process_chroma_select_simple(question=question,k=3,bot=bot)
        combine_question = tool.combine_vb_ans(contents=contents,question=question)
        cm_sc = [ i+"**"+j for i, j in zip(come,score)]
        # comes = "||".join(cm_sc)
        comes = cm_sc
        print(f"combine_question:{combine_question}")
        logging.info(f"combine_question:{combine_question}")
        results = deepseek_stream(combine_question)
        for item in results:
            if item.choices[0].finish_reason == "stop":
                item.choices[0].delta.content = comes
            yield item
        print("ds_Template_Warehouse-pass---------------------------------------------")

    # deepseek_内部资料模块:
    elif bot == "Company_Base":
        print('内部资料模块')
        logging.info('内部资料模块')
        #普通版
        # contents, comes = tool.process_chroma_select_simple(question=question, k=10, bot=bot)
        # combine_question = tool.combine_vb_ans(contents=contents,question=question)
        # print(f"combine_question:{combine_question}")

        # 判断db检索内容是否和问题相关并行版
        contents, comes = tool.process_chroma_select_mutil(question=question, k=100, bot=bot,function=deepseek_chat)
        if len(contents) == 0:
            yield '根据您的问题，暂时没有资料能匹配回答。'
        else:
            combine_question = tool.combine_vb_ans(contents=contents,question=question)
            print(f"combine_question:{combine_question}")
            logging.info(f"combine_question:{combine_question}")
            results = deepseek_stream(combine_question)
            for item in results:
                if item.choices[0].finish_reason == "stop":
                    item.choices[0].delta.content = comes
                yield item
        print("ds_Company_Base--------------------------------------------")

    #doubao_行业研报模块
    elif bot == "out_search":
        print(f"行业研报模块")
        logging.info(f"行业研报模块")
        qts = api.chat_out_search(question)
        print(f"拆分后的问题：{qts}")
        logging.info(f"拆分后的问题：{qts}")
        come = []
        contents = []
        cqs = []
        for idx, qt in enumerate(qts):
            #
            # ct,cm,sc = tool.process_chroma_select_simple(question=qt,k=20,bot=bot)
            ct,cm = tool.process_chroma_select_mutil_v1(question=qt, k=100, bot=bot,function=doubao_chat)
            # come = come + [[i,j]for i, j in zip(cm,sc)]
            come = come + cm
            contents = contents + ct
            combine_question = tool.combine_vb_ans(contents=ct,question=qt)
            cqs.append(combine_question)
        new_come = [i.split("**") for i in come]
        filter_come =  tool.remove_duplicate_rows(new_come)
        # cm_sc = [i[0] + "**" + i[1] for i in filter_come]
        # comes = [cm_sc]
        # comes = cm_sc
        comes = [i[0]+"**"+i[1] for i in filter_come]
        ans = mutil_function(doubao_chat, cqs)
        useful_ans = [i for i in ans]
        out_puts = api.combine_ans_out_search(ans=useful_ans, questions=qts, o_q=question)
        for item in out_puts:
            if item.choices[0].finish_reason == "stop":
                item.choices[0].delta.content = comes
            yield item
        print("豆包-out_search-pass---------------------------------------------")

    else:
        results = doubao_chat_stream(question)
        for item in results:
            yield item



def ten_q():
    from datetime import datetime
    import pandas as pd
    questions = ["顺丰时效件收入增速", "顺丰同城发展历程", "中通单票运输成本", "通达系转运中心数量", "顺丰与通达系利润",
                 "跨境电商市场规模", "国内企业海外仓数量", "行业电商件量占比", "网购渗透率", "零担快运市场规模"]
    questions = ["集团公司在规范化管理方面的规范性文件覆盖率是多少？",
                 "各单位在推进规范化管理体系建设时，应如何进行流程建设与管理？",
                 "2024年寄递业务中，特快业务在农村市场有哪些具体的发展策略？",
                 "2024年寄递业务在提升服务质量方面，对客服能力建设有哪些具体举措？",
                 "2024年快包业务在强化能力投入提升作业效率方面有哪些措施？",
                 "2024年寄递业务国际业务在运营机制和关键环节提升方面有哪些工作重点？",
                 "2024年寄递业务在深化陆运网改革方面，特快 “专网” 组网有哪些具体要求？",
                 "从机构角度看，2024年前三季度邮政集团在业务发展方面存在哪些不平衡现象？",
                 "在2024年第8次生产经营分析会上，对寄递业务四季度工作提出了哪些具体要求？",
                 "2024年中国邮政在服务国家重大战略方面取得了哪些成果？"]

if __name__ == '__main__':
    print("ok")
    # ten_q()
    # question = ["顺丰同城发展历程?","通达系转运中心数量?","顺丰与通达系利润?"]
    # api = Api()
    # for q in question:
    #     qts = api.chat_out_search(q)
    #     print(f"q:{q}, ans:{qts}")

    # print(Tools.classify_question("顺丰2024年营业收入是多少？"))
    for i in range(1):
        question = "顺丰2024年利润是多少？"
        data = RequestData(
            messages=[MessageItem(role='user', content=question)],
            chatId='45566',
            bot="out_search",
            temperature=0.1,
            stream=1
        )
        txt = ''
        for item in chat_base(data):
            if item.choices[0].finish_reason == "stop":
                print(item)
                break
            txt += item.choices[0].delta.content
        print("放回给用户的答案是：")
        print(txt)
        print("____________________________________________________")

    # #
    # question = "2025年经营发展目标是什么？"
    # data = RequestData(
    #     messages=[MessageItem(role='user', content=question)],
    #     chatId='45566',
    #     bot="Company_Base",
    #     temperature=0.1,
    #     stream=1
    # )
    # txt = ''
    # for item in chat_base(data):
    #     if item.choices[0].finish_reason == "stop":
    #         print(item)
    #         break
    #     txt += item.choices[0].delta.content
    #
    # print(txt)
    # #
    # # # question1 = "揽投部经理如何做好网格化管理？"
    # # # question2= "揽投部经理的岗位职责是什么？"
    # # # question3 = "揽投部经理如何对员工进行心理疏导？"
    # question = "2023年京东的現金流入是多少"
    # data = RequestData(
    #     messages=[MessageItem(role='user', content=question)],
    #     chatId='45566',
    #     bot="Template_Warehouse",
    #     temperature=0.1,
    #     stream=1
    # )
    # txt = ''
    # for item in chat_base(data):
    #     if item.choices[0].finish_reason == "stop":
    #         print(item)
    #         break
    #     txt += item.choices[0].delta.content
    #
    # print(txt)






